use std::num::NonZeroU128;

use solutiont8::{millar_rabin, MILLAR_RABIN_WITNESSES};

const MAX_CONSECUTIVE_PRIME_TEST: u32 = 100;

pub fn find_max_prime_factor(mut number: u128) -> u128 {
    fn isqrt(n: u128) -> u64 {
        (n as f64).sqrt() as u64
    }

    if millar_rabin(number) {
        return number;
    }

    use solutiont1::primes::primes;
    let mut primes = primes().with_limit(isqrt(number));
    let mut not_factor = 0;
    while let Some(p) = primes.next().map(u128::from) {
        if number % p == 0 {
            while number % p == 0 {
                number /= p;
            }
            primes.set_limit(isqrt(number));
            if number == 1 {
                return p;
            }
        } else {
            not_factor += 1;
            if not_factor == MAX_CONSECUTIVE_PRIME_TEST {
                if millar_rabin(number) {
                    return number;
                } else if let Some(f) = pollard_rho(number) {
                    number /= f.get();
                }
            }
        }
    }
    number
}

fn pollard_rho(n: u128) -> Option<NonZeroU128> {
    fn f(x: u128, n: u128) -> u128 {
        (x * x + 1) % n
    }
    fn gcd(a: u128, b: u128) -> u128 {
        match b {
            0 => a,
            _ => gcd(b, a % b),
        }
    }

    for &x in MILLAR_RABIN_WITNESSES {
        let mut x = x;
        let mut y = x;
        let mut d = 1;
        for _ in 0..100 {
            x = f(x, n);
            y = f(f(y, n), n);
            d = gcd(x.abs_diff(y), n);
            if d != 1 {
                break;
            }
        }
        if d != n {
            return NonZeroU128::new(d);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;

    #[track_caller]
    fn test_case(num: u128, factor: u128) {
        assert_eq!(find_max_prime_factor(num), factor);
    }

    #[test]
    #[ignore]
    fn test_big_number() {
        test_case(150382866733217945772777392401, 387792298444951);
        test_case(79004820088962844699776756200627, 203729729563409477);
    }
}
